﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Data;
using System.Windows.Documents;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Navigation;
using System.Windows.Shapes;
using Microsoft.Kinect;
using System.IO;

namespace KinectWeatherMap
{
    /// <summary>
    /// Interaction logic for MainWindow.xaml
    /// </summary>
    public partial class MainWindow : Window
    {
        enum PoseType
        {
            PoseLeft,
            PoseRight,
            Unknown
        }

        PoseType lastPoseRecognized = PoseType.Unknown;

        DepthImagePixel[] depthData;
        byte[] rgbBytes;
        Skeleton[] skeletons;
        ColorImagePoint[] colorCoordinates;

        int currentBackground = -1;
        List<Uri> bgImages = new List<Uri>();

        KinectSensor kinect;

        WriteableBitmap writableBitmap;
        byte[] foregroundImage;

        bool[] foregroundMask;

        bool growForegroundMask;

        public MainWindow()
        {
            InitializeComponent();

            bgImages.Add(new Uri("http://radar.weather.gov/ridge/Conus/Loop/NatLoop_Small.gif"));
            bgImages.Add(new Uri("http://radar.weather.gov/ridge/Conus/Loop/southeast_loop.gif"));
            bgImages.Add(new Uri("http://radar.weather.gov/ridge/Conus/Loop/northeast_loop.gif"));
            bgImages.Add(new Uri("http://radar.weather.gov/ridge/Conus/Loop/pacnorthwest_loop.gif"));
            bgImages.Add(new Uri("http://radar.weather.gov/ridge/Conus/Loop/pacsouthwest_loop.gif"));

            //backup images in case internet weather is not available
            //bgImages.Add(new Uri(System.IO.Path.GetFullPath("Images/NatLoop_Small.gif")));
            //bgImages.Add(new Uri(System.IO.Path.GetFullPath("Images/southeast_loop.gif")));
            //bgImages.Add(new Uri(System.IO.Path.GetFullPath("Images/northeast_loop.gif")));
            //bgImages.Add(new Uri(System.IO.Path.GetFullPath("Images/pacnorthwest_loop.gif")));
            //bgImages.Add(new Uri(System.IO.Path.GetFullPath("Images/pacsouthwest_loop.gif")));

            CycleBackground();

            this.Loaded += (s, e) =>
            {
                var sensor = KinectSensor.KinectSensors.FirstOrDefault();
                InitKinect(sensor);
                KinectSensor.KinectSensors.StatusChanged += Kinects_StatusChanged;
            };

            this.Closed += (s, e) =>
                {
                    UnsubscribeKinect(this.kinect);
                };

            this.KeyDown += new KeyEventHandler(MainWindow_KeyDown);
        }

        void MainWindow_KeyDown(object sender, KeyEventArgs e)
        {
            if (e.Key == Key.M)
            {
                growForegroundMask = !growForegroundMask;
            }
        }

        void Kinects_StatusChanged(object sender, StatusChangedEventArgs e)
        {
            switch (e.Status)
            {
                case KinectStatus.Connected:
                    if (this.kinect == null)
                    {
                        InitKinect(e.Sensor);
                    }
                    break;
                case KinectStatus.Disconnected:
                    if (e.Sensor == kinect)
                    {
                        UnsubscribeKinect(e.Sensor);
                    }
                    break;
            }
        }

        void UnsubscribeKinect(KinectSensor sensor)
        {
            if (sensor == null)
                return;
            sensor.AllFramesReady -= kinect_AllFramesReady;
            //sensor.Stop();
            this.kinect = null;
        }

        void InitKinect(KinectSensor sensor)
        {
            if (sensor == null)
                return;
            this.kinect = sensor;
            kinect.Start();

            kinect.DepthStream.Enable(DepthImageFormat.Resolution640x480Fps30);
            kinect.ColorStream.Enable(ColorImageFormat.RgbResolution640x480Fps30);
            kinect.SkeletonStream.Enable();

            kinect.AllFramesReady += kinect_AllFramesReady;
            //kinect.ElevationAngle = -7;
        }

        void kinect_AllFramesReady(object sender, AllFramesReadyEventArgs e)
        {
            using (var rgbFrame = e.OpenColorImageFrame())
            {
                if (rgbFrame == null)
                    return;

                if (rgbBytes == null || rgbBytes.Length != rgbFrame.PixelDataLength)
                {
                    rgbBytes = new byte[rgbFrame.PixelDataLength];
                }
                rgbFrame.CopyPixelDataTo(rgbBytes);
            }

            using (var depthFrame = e.OpenDepthImageFrame())
            {
                if (depthFrame == null)
                    return;

                if (depthData == null || depthData.Length != depthFrame.PixelDataLength)
                {
                    depthData = new DepthImagePixel[depthFrame.PixelDataLength];
                }
                depthFrame.CopyDepthImagePixelDataTo(depthData);
            }

            //ProcessRGBImage();
            //ProcessRGBWithBackgroundRemoval();
            ProcessRGBWithBackgroundRemovalAligned();

            using (var skeletonFrame = e.OpenSkeletonFrame())
            {
                if (skeletonFrame == null)
                    return;

                if (skeletons == null || skeletons.Length != skeletonFrame.SkeletonArrayLength)
                {
                    skeletons = new Skeleton[skeletonFrame.SkeletonArrayLength];
                }
                skeletonFrame.CopySkeletonDataTo(skeletons);
            }

            ProcessSkeletons();
        }

        private void ProcessRGBImage()
        {
            if (kinect == null)
                return;

            int rgbWidth = kinect.ColorStream.FrameWidth;
            int rgbHeight = kinect.ColorStream.FrameHeight;

            int numPixels = rgbWidth * rgbHeight;
            for (int i = 0; i < numPixels; i++)
            {
                rgbBytes[i * 4 + 3] = 255;
            }

            if (writableBitmap == null ||
                writableBitmap.Width != rgbWidth ||
                writableBitmap.Height != rgbHeight)
            {
                writableBitmap =
                    new WriteableBitmap(rgbWidth, rgbHeight,
                                        96, 96,
                                        PixelFormats.Bgra32, null);

                kinectImage.Source = writableBitmap;
            }
            writableBitmap.WritePixels(
                new Int32Rect(0, 0, rgbWidth, rgbHeight),
                rgbBytes,
                rgbWidth * 4,
                0);
        }

        private void ProcessRGBWithBackgroundRemoval()
        {
            if (kinect == null)
                return;

            int rgbWidth = kinect.ColorStream.FrameWidth;
            int rgbHeight = kinect.ColorStream.FrameHeight;
            int depthWidth = kinect.DepthStream.FrameWidth;
            int depthHeight = kinect.DepthStream.FrameHeight;

            if (foregroundImage == null ||
                foregroundImage.Length != depthWidth * depthHeight * 4)
            {
                foregroundImage = new byte[depthWidth * depthHeight * 4];
            }

            Array.Clear(foregroundImage, 0, foregroundImage.Length);

            for (int y = 0; y < depthHeight; y++)
            {
                for (int x = 0; x < depthWidth; x++)
                {
                    int depthIndex = (x + y * depthWidth);

                    //We don't need it here, but this is how you would get the depth values
                    //short depth = depthData[targetIndex].Depth;

                    int playerIndex = depthData[depthIndex].PlayerIndex;
                    if (playerIndex == 0)
                    {
                        continue;
                    }

                    int colorIndex = depthIndex * 4;

                    foregroundImage[depthIndex * 4 + 0] = rgbBytes[colorIndex + 0];
                    foregroundImage[depthIndex * 4 + 1] = rgbBytes[colorIndex + 1];
                    foregroundImage[depthIndex * 4 + 2] = rgbBytes[colorIndex + 2];
                    foregroundImage[depthIndex * 4 + 3] = 255;
                }
            }

            if (writableBitmap == null ||
                writableBitmap.Width != depthWidth ||
                writableBitmap.Height != depthHeight)
            {
                writableBitmap = new WriteableBitmap(depthWidth, depthHeight,
                                                        96, 96,
                                                        PixelFormats.Bgra32, null);

                kinectImage.Source = writableBitmap;
            }

            writableBitmap.WritePixels(new Int32Rect(0, 0, depthWidth, depthHeight),
                                        foregroundImage, depthWidth * 4, 0);
        }

        private void ProcessRGBWithBackgroundRemovalAligned()
        {
            if (kinect == null)
                return;

            int rgbWidth = kinect.ColorStream.FrameWidth;
            int rgbHeight = kinect.ColorStream.FrameHeight;
            int depthWidth = kinect.DepthStream.FrameWidth;
            int depthHeight = kinect.DepthStream.FrameHeight;

            if (foregroundImage == null ||
                foregroundImage.Length != depthWidth * depthHeight * 4)
            {
                foregroundImage = new byte[depthWidth * depthHeight * 4];
            }

            if (colorCoordinates == null || colorCoordinates.Length != depthData.Length)
            {
                colorCoordinates = new ColorImagePoint[depthData.Length];
            }

            CoordinateMapper mapper = kinect.CoordinateMapper;

            mapper.MapDepthFrameToColorFrame(kinect.DepthStream.Format, depthData, kinect.ColorStream.Format, colorCoordinates);

            Array.Clear(foregroundImage, 0, foregroundImage.Length);

            for (int y = 0; y < depthHeight; y++)
            {
                for (int x = 0; x < depthWidth; x++)
                {
                    int depthIndex = (x + y * depthWidth);

                    //We don't need it here, but this is how you would get the depth values
                    //short depth = depthData[targetIndex].Depth;

                    int playerIndex = depthData[depthIndex].PlayerIndex;
                    if (playerIndex == 0)
                    {
                        continue;
                    }

                    var coord = colorCoordinates[depthIndex];
                    int colorX = coord.X;
                    int colorY = coord.Y;

                    if (colorX < 0 || colorX >= rgbWidth ||
                        colorY < 0 || colorY >= rgbHeight)
                        continue;

                    int colorIndex = 4 * (colorX + colorY * rgbWidth);

                    foregroundImage[depthIndex * 4 + 0] = rgbBytes[colorIndex + 0];
                    foregroundImage[depthIndex * 4 + 1] = rgbBytes[colorIndex + 1];
                    foregroundImage[depthIndex * 4 + 2] = rgbBytes[colorIndex + 2];
                    foregroundImage[depthIndex * 4 + 3] = 255;
                }
            }

            if (writableBitmap == null ||
                writableBitmap.Width != depthWidth ||
                writableBitmap.Height != depthHeight)
            {
                writableBitmap = new WriteableBitmap(depthWidth, depthHeight,
                                                        96, 96,
                                                        PixelFormats.Bgra32, null);

                kinectImage.Source = writableBitmap;
            }

            writableBitmap.WritePixels(new Int32Rect(0, 0, depthWidth, depthHeight),
                                        foregroundImage, depthWidth * 4, 0);
        }

        #region Other Approaches

        private void ProcessRGBWithBackgroundRemovalColorSpace()
        {
            if (kinect == null)
                return;

            int rgbWidth = kinect.ColorStream.FrameWidth;
            int rgbHeight = kinect.ColorStream.FrameHeight;
            int depthWidth = kinect.DepthStream.FrameWidth;
            int depthHeight = kinect.DepthStream.FrameHeight;

            if (foregroundImage == null ||
                foregroundImage.Length != rgbWidth * rgbHeight * 4)
            {
                foregroundImage = new byte[rgbWidth * rgbHeight * 4];
            }

            if (colorCoordinates == null || colorCoordinates.Length != depthData.Length)
            {
                colorCoordinates = new ColorImagePoint[depthData.Length];
            }
            DepthImagePoint[] depthMapping = new DepthImagePoint[rgbWidth * rgbHeight];

            CoordinateMapper mapper = kinect.CoordinateMapper;
            mapper.MapColorFrameToDepthFrame(kinect.ColorStream.Format, kinect.DepthStream.Format, depthData, depthMapping);
            //mapper.MapDepthFrameToColorFrame(kinect.DepthStream.Format, depthData, kinect.ColorStream.Format, colorCoordinates);

            Array.Clear(foregroundImage, 0, foregroundImage.Length);

            for (int y = 0; y < rgbHeight; y++)
            {
                for (int x = 0; x < rgbWidth; x++)
                {
                    int targetIndex = (x + y * depthWidth);

                    //We don't need it here, but this is how you would get the depth values
                    //short depth = depthData[targetIndex].Depth;

                    var depthPoint = depthMapping[targetIndex];

                    if (depthPoint.X < 0 || depthPoint.X >= depthWidth ||
                        depthPoint.Y < 0 || depthPoint.Y >= depthHeight)
                        continue;

                    int depthIndex = depthPoint.X + depthPoint.Y * depthWidth;

                    int playerIndex = depthData[depthIndex].PlayerIndex;
                    if (playerIndex == 0)
                    {
                        continue;
                    }

                    foregroundImage[targetIndex * 4 + 0] = rgbBytes[targetIndex * 4 + 0];
                    foregroundImage[targetIndex * 4 + 1] = rgbBytes[targetIndex * 4 + 1];
                    foregroundImage[targetIndex * 4 + 2] = rgbBytes[targetIndex * 4 + 2];
                    foregroundImage[targetIndex * 4 + 3] = 255;
                }
            }

            if (writableBitmap == null ||
                writableBitmap.Width != depthWidth ||
                writableBitmap.Height != depthHeight)
            {
                writableBitmap = new WriteableBitmap(depthWidth, depthHeight,
                                                        96, 96,
                                                        PixelFormats.Bgra32, null);

                kinectImage.Source = writableBitmap;
            }

            writableBitmap.WritePixels(new Int32Rect(0, 0, depthWidth, depthHeight),
                                        foregroundImage, depthWidth * 4, 0);
        }

        private void ProcessRGBWithBackgroundRemovalGrowing()
        {
            //This approach isn't fully working
            if (kinect == null)
                return;

            int rgbWidth = kinect.ColorStream.FrameWidth;
            int rgbHeight = kinect.ColorStream.FrameHeight;
            int depthWidth = kinect.DepthStream.FrameWidth;
            int depthHeight = kinect.DepthStream.FrameHeight;

            if (foregroundImage == null ||
                foregroundImage.Length != depthData.Length * 4)
            {
                foregroundImage = new byte[depthData.Length * 4];
            }

            if (colorCoordinates == null || colorCoordinates.Length != depthData.Length)
            {
                colorCoordinates = new ColorImagePoint[depthData.Length];
            }

            if (foregroundMask == null ||
                foregroundMask.Length != depthData.Length)
            {
                foregroundMask = new bool[depthData.Length];
            }

            CoordinateMapper mapper = kinect.CoordinateMapper;

            mapper.MapDepthFrameToColorFrame(kinect.DepthStream.Format, depthData, kinect.ColorStream.Format, colorCoordinates);


            UpdateForegroundMaskFromPlayerMask();

            for (int i = 0; i < 1; i++)
                GrowForegroundMask();

            UpdateForegroundImageFromForegroundMask();

            if (writableBitmap == null ||
                writableBitmap.Width != depthWidth ||
                writableBitmap.Height != depthHeight)
            {
                writableBitmap = new WriteableBitmap(depthWidth, depthHeight,
                                                        96, 96,
                                                        PixelFormats.Bgra32, null);

                kinectImage.Source = writableBitmap;
            }

            writableBitmap.WritePixels(new Int32Rect(0, 0, depthWidth, depthHeight),
                                        foregroundImage, depthWidth * 4, 0);
        }

        private void GrowForegroundMask()
        {
            int depthWidth = kinect.DepthStream.FrameWidth;
            int depthHeight = kinect.DepthStream.FrameHeight;

            int unknownDepth = kinect.DepthStream.UnknownDepth+1;
            int tooNearDepth = kinect.DepthStream.TooNearDepth;

            for (int y = 1; y < depthHeight - 1; y++)
            {
                for (int x = 1; x < depthWidth - 1; x++)
                {
                    int index = (x + y * depthWidth);

                    int currentDepth = depthData[index].Depth;

                    //if (x == 320 && y == 240)
                    //{
                    //    txtPos.Text = currentDepth.ToString();
                    //}

                    if (foregroundMask[index] == true)
                    {
                        continue;
                    }
                    

                    int leftDepth = -5;
                    int rightDepth = -5;
                    int upDepth = -5;
                    int downDepth = -5;
                    bool leftForeground = false;
                    bool rightForeground = false;
                    bool upForeground = false;
                    bool downForeground = false;

                    int threshold = 100;

                    leftForeground = foregroundMask[index - 1];
                    leftDepth = depthData[index - 1].Depth;

                    if (leftForeground &&
                        (Math.Abs(leftDepth - currentDepth) < threshold ||
                         currentDepth == tooNearDepth ||
                         currentDepth == unknownDepth))
                    {
                        foregroundMask[index] = true;
                        continue;
                    }

                    rightForeground = foregroundMask[index + 1];
                    rightDepth = depthData[index + 1].Depth;

                    if (rightForeground &&
                        (Math.Abs(rightDepth - currentDepth) < threshold ||
                         currentDepth == tooNearDepth ||
                         currentDepth == unknownDepth))
                    {
                        foregroundMask[index] = true;
                        continue;
                    }

                    upForeground = foregroundMask[index - depthWidth];
                    upDepth = depthData[index - depthWidth].Depth;

                    if (upForeground &&
                        (Math.Abs(upDepth - currentDepth) < threshold ||
                         currentDepth == tooNearDepth ||
                         currentDepth == unknownDepth))
                    {
                        foregroundMask[index] = true;
                        continue;
                    }

                    downForeground = foregroundMask[index + depthWidth];
                    downDepth = depthData[index + depthWidth].Depth;

                    if (downForeground &&
                        (Math.Abs(downDepth - currentDepth) < threshold ||
                         currentDepth == tooNearDepth ||
                         currentDepth == unknownDepth))
                    {
                        foregroundMask[index] = true;
                        continue;
                    }
                }
            }
        }

        private void UpdateForegroundMaskFromPlayerMask()
        {
            int depthWidth = kinect.DepthStream.FrameWidth;
            int depthHeight = kinect.DepthStream.FrameHeight;

            Array.Clear(foregroundMask, 0, foregroundMask.Length);

            for (int y = 0; y < depthHeight; y++)
            {
                for (int x = 0; x < depthWidth; x++)
                {
                    int index = (x + y * depthWidth);

                    //We don't need it here, but this is how you would get the depth values
                    //short depth = depthData[targetIndex].Depth;

                    int playerIndex = depthData[index].PlayerIndex;
                    if (playerIndex == 0)
                    {
                        continue;
                    }

                    foregroundMask[index] = true;
                }
            }
        }

        private void UpdateForegroundImageFromForegroundMask()
        {
            int rgbWidth = kinect.ColorStream.FrameWidth;
            int rgbHeight = kinect.ColorStream.FrameHeight;
            int depthWidth = kinect.DepthStream.FrameWidth;
            int depthHeight = kinect.DepthStream.FrameHeight;

            Array.Clear(foregroundImage, 0, foregroundImage.Length);

            for (int y = 0; y < depthHeight; y++)
            {
                for (int x = 0; x < depthWidth; x++)
                {
                    int index = (x + y * depthWidth);
                    int index4 = index * 4;

                    //if (x == 320 && y == 240)
                    //{
                    //    foregroundImage[index4 + 0] = 0;
                    //    foregroundImage[index4 + 1] = 0;
                    //    foregroundImage[index4 + 2] = 0;
                    //    foregroundImage[index4 + 3] = 255;
                    //    continue;
                    //}

                    bool isForeground = foregroundMask[index];

                    if (isForeground)
                    {
                        var coord = colorCoordinates[index];
                        int colorX = coord.X;
                        int colorY = coord.Y;

                        if (colorX < 0 || colorX >= rgbWidth ||
                            colorY < 0 || colorY >= rgbHeight)
                            continue;

                        int colorIndex = 4 * (colorX + colorY * rgbWidth);

                        foregroundImage[index4 + 0] = rgbBytes[colorIndex + 0];
                        foregroundImage[index4 + 1] = rgbBytes[colorIndex + 1];
                        foregroundImage[index4 + 2] = rgbBytes[colorIndex + 2];
                        foregroundImage[index4 + 3] = 255;
                    }
                }
            }
        }
        
        #endregion

        void ProcessSkeletons()
        {
            var skeleton = skeletons.Where(s => s.TrackingState == SkeletonTrackingState.Tracked)
                                    .OrderBy((s) => s.Position.Z * Math.Abs(s.Position.X))
                                    .FirstOrDefault();

            if (skeleton == null)
                return;

            var shoulderLeft = skeleton.Joints[JointType.ShoulderLeft];
            var shoulderRight = skeleton.Joints[JointType.ShoulderRight];

            RecognizePoses(skeleton.Position, shoulderLeft, shoulderRight);
        }

        private void RecognizePoses(SkeletonPoint skeletonPosition, Joint shoulderLeft, Joint shoulderRight)
        {
            if (shoulderLeft.TrackingState != JointTrackingState.Tracked ||
                shoulderRight.TrackingState != JointTrackingState.Tracked)
            {
                return;
            }

            //Test for left pose
            if (skeletonPosition.X < -0.15 &&
                shoulderRight.Position.Z > shoulderLeft.Position.Z + 0.15)
            {
                if (lastPoseRecognized == PoseType.PoseRight)
                {
                    CycleBackground();
                }
                lastPoseRecognized = PoseType.PoseLeft;
            }
            //Test for right pose
            else if (skeletonPosition.X > 0.15 &&
                     shoulderLeft.Position.Z > shoulderRight.Position.Z + 0.15)
            {
                if (lastPoseRecognized == PoseType.PoseLeft)
                {
                    CycleBackground();
                }
                lastPoseRecognized = PoseType.PoseRight;
            }
        }

        void CycleBackground()
        {
            currentBackground++;
            if (currentBackground > bgImages.Count - 1)
            {
                currentBackground = 0;
            }
            weatherImage.GifSource = bgImages[currentBackground];
        }
    }
}
